/*
 * FfmpegTsMuxer.cpp
 *
 *  Created on: 2016年8月24日
 *      Author: zhengboyuan
 */

#include "FfmpegTsMuxer.h"
#include "CLog.h"
#include "AacHelper.h"


namespace av
{


FfmpegTsMuxer::FfmpegTsMuxer() :
	m_sink(),
	m_fmtCtx(),
	m_videoIndex(-1),
	m_audioIndex(-1),
	m_pts()
{
}

FfmpegTsMuxer::~FfmpegTsMuxer()
{
}

int FfmpegTsMuxer::open(const MediaFormat& fmt)
{
	m_fmt = fmt;

	int rc = avformat_alloc_output_context2(&m_fmtCtx, NULL, "mpegts", NULL);
	if (!m_fmtCtx)
	{
		return rc;
	}

	if (m_fmt.hasVideo())
	{
		AVCodecID codecID = (AVCodecID)m_fmt.m_codec;
		if (m_fmt.m_codec == MEDIA_CODEC_HEVC)
		{
			codecID = AV_CODEC_ID_HEVC;
		}

		AVCodec* codec = avcodec_find_decoder(codecID);
		if (codec != NULL)
		{
			AVStream *out_stream = avformat_new_stream(m_fmtCtx, NULL);
			AVCodecContext * codecCtx = out_stream->codec;

			avcodec_get_context_defaults3(codecCtx, codec);

			if (m_fmt.m_width <= 0)
			{
				codecCtx->width = 1920;
				codecCtx->height = 1080;
			}
			else
			{
				codecCtx->width = m_fmt.m_width;
				codecCtx->height = m_fmt.m_height;
				codecCtx->profile = m_fmt.m_profile;
			}

			codecCtx->time_base = av_make_q(1, 90000);

			out_stream->time_base = codecCtx->time_base;

			// extra_data
			if (m_fmt.m_videoProp.size() > 0)
			{
				codecCtx->extradata = (uint8_t*)av_memdup(m_fmt.m_videoProp.c_str(), m_fmt.m_videoProp.size());
				codecCtx->extradata_size = m_fmt.m_videoProp.size();
			}

			if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
			{
				out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
			}

			m_videoIndex = 0;
		}
		else
		{
			CLog::warning("no such video codec. id:%d\n", m_fmt.m_codec);
		}
	}

	if (m_fmt.hasAudio())
	{
		AVCodecID codecID = (AVCodecID)m_fmt.m_audioCodec;
		AVCodec* codec = avcodec_find_decoder(codecID);
		if (codec != NULL)
		{
			AVStream *out_stream = avformat_new_stream(m_fmtCtx, codec);
			AVCodecContext * codecCtx = out_stream->codec;

			avcodec_get_context_defaults3(codecCtx, codec);

			codecCtx->channels = m_fmt.m_channels;
			codecCtx->channel_layout = av_get_default_channel_layout(m_fmt.m_channels);
			codecCtx->sample_rate = m_fmt.m_sampleRate;
			codecCtx->time_base.num = 1;
			codecCtx->time_base.den = 90000;

			out_stream->time_base = codecCtx->time_base;

			if (m_fmt.m_audioConfig.size() > 0)
			{
				codecCtx->extradata = (uint8_t*)av_memdup(m_fmt.m_audioConfig.c_str(), m_fmt.m_audioConfig.size());
				codecCtx->extradata_size = m_fmt.m_audioConfig.size();
			}

			if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
			{
				out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
			}

			m_audioIndex = m_videoIndex + 1;
		}
		else
		{
			CLog::warning("no such audio codec. id:%d\n", m_fmt.m_codec);
		}
	}

	if ((m_audioIndex < 0) && (m_videoIndex < 0))
	{
		avformat_free_context(m_fmtCtx);
		m_fmtCtx = NULL;
		return ENODEV;
	}

	m_fmtCtx->pb = createIOContext();
	if (m_fmtCtx->pb == NULL)
	{
		avformat_free_context(m_fmtCtx);
		m_fmtCtx = NULL;
		m_videoIndex = -1;
		m_audioIndex = -1;
		return ENODEV;
	}

	AVDictionary *options = NULL;
	//av_dict_set(&options, "pat_period", "10", 0);
	//av_dict_set(&options, "sdt_period", "10", 0);
	//av_dict_set(&options, "pcr_period", "5", 0);
	//av_dict_set(&options, "omit_video_pes_length", "0", 0);
	//av_dict_set(&options, "mpegts_pmt_start_pid", "0x150", 0);
	
	rc = avformat_write_header(m_fmtCtx, &options);
	
	av_dict_free(&options);

	return rc;
}

void FfmpegTsMuxer::close()
{
	if (m_fmtCtx)
	{
		av_write_trailer(m_fmtCtx);

		//avio_close(m_fmtCtx->pb);

		avformat_free_context(m_fmtCtx);
		m_fmtCtx = NULL;
	}

	m_videoIndex = -1;
	m_audioIndex = -1;
}

bool FfmpegTsMuxer::isOpen()
{
	return (m_fmtCtx != NULL);
}

bool FfmpegTsMuxer::write(int type, uint8_t* data, int size, int64_t pts, int duration, int flags)
{
	if (!isOpen())
	{
		return false;
	}

	AVPacket pkt;
	av_init_packet(&pkt);
	pkt.data = data;
	pkt.size = size;
	pkt.pts = pts;
	pkt.dts = pts;
	pkt.flags = flags;

	if (type == MEDIA_TYPE_VIDEO)
	{
		pkt.stream_index = m_videoIndex;
	}
	else if (type == MEDIA_TYPE_AUDIO)
	{
		pkt.stream_index = m_audioIndex;
	}
	
	if (pkt.stream_index < 0)
	{
		return false;
	}

	int clockRate = AV_TIME_BASE;
	if (pkt.stream_index == m_videoIndex)
	{
		clockRate = m_fmt.m_clockRate > 0 ? m_fmt.m_clockRate : AV_TIME_BASE;
	}
	else
	{
		clockRate = m_fmt.m_audioRate > 0 ? m_fmt.m_audioRate : m_fmt.m_sampleRate;

		if (m_fmt.m_audioCodec == MEDIA_CODEC_AAC)
		{
			if (!AacHelper::isAdts(pkt.data, pkt.size))
			{
				m_buffer.ensure(pkt.size + AacHelper::ADTS_HEADER_SIZE * 2);
				AacHelper::makeAdtsHeader(m_fmt.m_audioProfile, m_fmt.m_sampleRate, m_fmt.m_channels,
					pkt.size + AacHelper::ADTS_HEADER_SIZE,
					m_buffer.data(), m_buffer.capacity());
				memcpy(m_buffer.data() + AacHelper::ADTS_HEADER_SIZE, pkt.data, pkt.size);
				pkt.data = m_buffer.data();
				pkt.size = pkt.size + AacHelper::ADTS_HEADER_SIZE;
			}
		}
	}

	AVRational src = av_make_q(1, clockRate);
	AVRational dst = av_make_q(1, 90000);
	av_packet_rescale_ts(&pkt, src, dst);

	m_pts = pkt.pts;

	int ret = av_interleaved_write_frame(m_fmtCtx, &pkt);

	return ret == 0;
}

void FfmpegTsMuxer::setSink(TsMuxerSink* sink)
{
	m_sink = sink;
}


AVIOContext* FfmpegTsMuxer::createIOContext()
{
    int buffer_size = 188 * 7;
    unsigned char *buffer = (unsigned char *)av_malloc(buffer_size);

    return avio_alloc_context(buffer, buffer_size, 1, this,
        NULL, ts_write_packet, NULL);
}

int FfmpegTsMuxer::ts_write_packet(void *opaque, uint8_t *buf, int buf_size)
{
	FfmpegTsMuxer* pthis = (FfmpegTsMuxer*)opaque;
	if (pthis)
	{


		pthis->writePacket(buf, buf_size);
	}
	return buf_size;
}

void FfmpegTsMuxer::writePacket(uint8_t *buf, int buf_size)
{
	if (m_sink)
	{
		m_sink->writePacket(buf, buf_size, m_pts, 0);
	}
}



} /* namespace av */
